import os
from traceback import print_tb
from jax import grad
import jax.numpy as jnp
from jax import jit
import time
import matplotlib
import numpy as np
import numpy.random as npr
import jax
import jax.numpy as jnp
from jax import device_put
from jax import jit, grad, lax, random
from jax.example_libraries import optimizers
from jax.example_libraries import stax
from jax.example_libraries.stax import Dense, FanOut, Relu, Softplus, Sigmoid, FanInSum
from jax.nn import sigmoid
from functools import partial
from jax import vmap
from flax import linen as nn
from flax.training import train_state
from flax import struct
from jax import lax

from jax import tree_util
from jax.tree_util import tree_structure
from jax.tree_util import tree_flatten, tree_unflatten

import jax.experimental.sparse as sparse

import optax
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from scipy.integrate import odeint

from matplotlib.animation import FuncAnimation


def spring_damper_ode_2d(state, t, k, b, m, k_couple):
    """
    二维阻尼弹簧振子系统的微分方程
    """
    x, y, vx, vy = state
    dxdt = vx
    dydt = vy
    # dvxdt = -(k/m)*x - (b/m)*vx + (k_couple/m)*y
    # dvydt = -(k/m)*y - (b/m)*vy + (k_couple/m)*x
    dvxdt = -(k/m)*x - (b/m)*vx + (k_couple/m)*y*(1-k_couple/m)*x
    dvydt = -(k/m)*y - (b/m)*vy + (k_couple/m)*x*y*np.log(vy*vx*vx*vy)
    return [dxdt, dydt, dvxdt, dvydt]

def plot_trajectory_2d(k, b, m, k_couple, X, V, t_max=500.0, dt=0.01, disturbance=True):
    """
    绘制二维阻尼弹簧振子系统的相空间轨迹
    """
    x0, y0 = X
    vx0, vy0 = V
    t = np.arange(0, t_max, dt)
    s0 = [x0, y0, vx0, vy0]
    s = np.zeros((len(t), len(s0)))
    s[0, :] = s0

    v_increase_step = 20  # 每隔v_increase_step步, 注入一次能量
    v_input_i = 0
    v_input_seq = [
        [0.1,0.7],
        [0.4,0.3],
        [-0.7,0.1],
        [0.3,0.4],
        [-0.5,0.5],
        [0.2,0.2],
        [0.8,-0.8],
        [-0.6,0.6],
        [0.9,-0.3],
    ]
    def v_input_func():
        nonlocal v_input_i
        v_input_i += 1
        if v_input_i % len(v_input_seq) == 0:
            v_input_i = 0
        return v_input_seq[v_input_i % len(v_input_seq)]
    
    for i in range(1, len(t)):
        tspan = [t[i-1], t[i]]
        s[i, :] = odeint(spring_damper_ode_2d, s[i-1,:], tspan, args=(k, b, m, k_couple))[-1,:]
        if disturbance and i >= 700 and i % v_increase_step == 0:
            s[i, 2:] += v_input_func()

    if disturbance:
        alpha = np.linspace(0.1, 1, len(t))
        for i in range(1, len(t)):
            plt.plot(s[i-1:i+1, 0], s[i-1:i+1, 1], 'b-', alpha=alpha[i])
    else:
        plt.plot(s[:, 0], s[:, 1], 'b-')

    plt.plot(s0[0], s0[1], 'r.', markersize=10)
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Phase Space Trajectory')
    plt.xlim([-3, 3])
    plt.ylim([-3, 3])

def main_2d():
    # 二维阻尼弹簧振子系统
    k = 1.0  # 弹簧刚度系数
    b = 0.2  # 阻尼系数
    m = 1.0  # 质量
    k_couple = 0.5

    # 绘制二维阻尼弹簧振子系统的相空间轨迹
    # X = [1.0, 1.0]  # 初始位置
    # V = [0.3, 0.0]  # 初始速度
    X = np.random.uniform(-1.0, 1.0, size=(2,))  # 初始位置
    V = np.random.uniform(-1.0, 1.0, size=(2,))  # 初始速度

    plot_trajectory_2d(k, b, m, k_couple, X, V)

    plt.show()
    

if __name__ == '__main__':
    main_2d()